import { GobangGameState } from "../gobang/GobangGameState";
import { MCTSGameMove } from "./MCTSGameMove";
import { GameResult, MCTSGameState } from "./MCTSGameState";
import { LevelType, MCTSNode } from "./MCTSNode";

export class MCTS {
    root: MCTSNode;
    N_SIMULATION_TIMES: number; // 博弈树的节点数

    constructor(gameState: MCTSGameState, nSimulationTimes: number) {
        this.root = new MCTSNode(LevelType.Computer, gameState, null);
        this.N_SIMULATION_TIMES = nSimulationTimes;
    }

    /**
     * 产生下一步的走法
     */
    decideMove(): MCTSGameMove | undefined {
        this.generateGameTree();
        const selectedChild = this.root.selectChildWithMaxUCB();
        return selectedChild?.gameState.returnLastMove();
    }

    /**
     * 按照指定的搜索次数构建游戏树
     */
    generateGameTree(): void {
        for (let i = 1; i <= this.N_SIMULATION_TIMES; i++) {
            this.performOnePass();
        }
    }

    /**
     * 执行一次搜索流程
     */
    performOnePass(): void {
        const selected = this.select(this.root);
        const child = this.expand(selected); // selected 为终局节点时返回 null
        let computerWon: boolean;
        if (child!= null) {
            computerWon = this.simulate(child);
        } else {
            computerWon = (selected.gameState.judgeLastMove() === GameResult.ComputerWon); // 终局节点
        }
        this.backPropagate(child || selected, computerWon); // child!= null? child : selected

    }

    /**
     * 依据最大最小算法，选择第一个尚存在未展开子节点的节点
     */
    select(root: MCTSNode): MCTSNode {
        return this.recurSelect(root);
    }

    /**
     * 扩展传入节点
     */
    expand(node: MCTSNode): MCTSNode | null {
        if (node.existsUnexpandedChild()) {
            const stateOfChild = node.gameState.expand(node.getChildrenGameStates());
            const typeOfChild = node.getLevelType() === LevelType.Computer? LevelType.Player : LevelType.Computer;
            const child = new MCTSNode(typeOfChild, stateOfChild, node);
            node.appendChild(child);
            return child;
        } else {
            return null;
        }
    }

    /**
     * 从当前游戏状态开始，使用随机方法模拟一次完整的游戏，告知游戏结果
     * 不会更新当前节点的状态
     */
    simulate(MCTSNode: MCTSNode): boolean {
        return MCTSNode.gameState.simulate();
    }

    /**
     * 从当前节点开始，反向传播模拟结果
     * 将更新此节点的状态
     */
    backPropagate(MCTSNode: MCTSNode, won: boolean): void {
        this.recurBackPropagate(MCTSNode, won);
    }

    // ========================================== 递归功能函数 ==========================================
    /**
     * 递归选择子节点
     */
    recurSelect(cur: MCTSNode): MCTSNode {
        if (cur.existsUnexpandedChild()) {
            return cur; // 存在未扩展子节点
        } else if (cur.isLeaf()) {
            return cur; // 终局节点
        } else {
            return this.recurSelect(cur.selectChildWithMaxUCB()!);
        }
    }

    /**
     * 递归反向传播
     */
    recurBackPropagate(cur: MCTSNode, won: boolean): void {
        cur.updateStatus(won);
        if (cur.parent!= null) {
            this.recurBackPropagate(cur.parent, won);
        }
    }
}